package cn.tk.config;

import org.apache.http.Header;
import org.apache.http.HttpHost;
import org.apache.http.HttpResponse;
import org.apache.http.impl.client.DefaultConnectionKeepAliveStrategy;
import org.apache.http.message.BasicHeader;
import org.apache.http.protocol.HttpContext;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.annotation.Order;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Created by denglw on 2020/11/6.<br/>
 * Desc: ES连接配置
 */
@Configuration
@Order(1)
public class ElasticConfig {

    /**
     * ES实例列表
     */
    @Value("${spring.elasticsearch.rest.hosts}")
    private String hosts;

    /**
     * ES实例端口
     */
    @Value("${spring.elasticsearch.rest.port:9200}")
    private int port;

    @Bean
    public RestHighLevelClient client() {
        String[] hostList = hosts.split(",");
        HttpHost[] httpHosts = new HttpHost[hostList.length];
        for (int i = 0; i < hostList.length; i++) {
            httpHosts[i] = new HttpHost(hostList[i], port);
        }
        RestClientBuilder builder = RestClient.builder(httpHosts);

        // 异步httpclient连接延时配置
        builder.setRequestConfigCallback(requestConfigBuilder -> {
            requestConfigBuilder.setConnectionRequestTimeout(1000);
            requestConfigBuilder.setConnectTimeout(2000);
            requestConfigBuilder.setSocketTimeout(60_000);
            return requestConfigBuilder;
        });
        // 异步httpclient连接数配置
        builder.setHttpClientConfigCallback(httpClientBuilder -> {
            httpClientBuilder.setMaxConnTotal(100);
            httpClientBuilder.setMaxConnPerRoute(50);
            List<Header> headers = new ArrayList<>();
            headers.add(new BasicHeader("Connection", "keep-alive"));
            headers.add(new BasicHeader("Keep-Alive", "720"));
            httpClientBuilder.setDefaultHeaders(headers);
            httpClientBuilder.setKeepAliveStrategy(CustomKeepAliveStrategy.INSTANCE);
            return httpClientBuilder;
        });
        return new RestHighLevelClient(builder);
    }

    static class CustomKeepAliveStrategy extends DefaultConnectionKeepAliveStrategy {

        public static final CustomKeepAliveStrategy INSTANCE = new CustomKeepAliveStrategy();

        private CustomKeepAliveStrategy() {
            super();
        }

        @Override
        public long getKeepAliveDuration(HttpResponse response, HttpContext context) {
            long keepAliveDuration = super.getKeepAliveDuration(response, context);
            // <0 为无限期keepalive
            // 将无限期替换成一个默认的时间
            if (keepAliveDuration < 0) {
                /*
                 * 最大keep alive的时间（分钟）
                 * 这里默认为10分钟(600S)，可以根据实际情况设置。可以观察客户端机器状态为TIME_WAIT的TCP连接数，如果太多，可以增大此值。
                 */
                long maxKeepAliveSeconds = 600;
                return TimeUnit.SECONDS.toMillis(maxKeepAliveSeconds);
            }
            return keepAliveDuration;
        }
    }
}
